from abc import ABC, abstractmethod
import numpy as np


class ParamsInitBase(ABC):
    def __init__(self):
        pass

    # 获得卷积层的参数
    @abstractmethod
    def getParams(self, filter_size, filter_num, channels):
        pass

    @abstractmethod
    def getAffineParams(self, x_shape, output_size):
        pass


class He(ParamsInitBase):
    """
        He 参数初始化方法
    """

    def __init__(self):
        pass

    def getAffineParams(self, x_shape, output_size):
        if len(x_shape) == 4:
            shape = x_shape[1]*x_shape[2]*x_shape[3]
        else:
            shape = x_shape[1]
        weight = np.math.sqrt(2 / shape)
        x = np.random.randn(shape, output_size) * weight
        b = np.zeros(output_size)
        return x, b

    def getParams(self, filter_size, filter_num, channels):
        weight = np.math.sqrt(2 / (filter_num * filter_size * filter_size))
        x = np.random.randn(filter_num, channels,
                            filter_size, filter_size) * weight
        b = np.zeros((filter_num))
        return x, b


class Xavier(ParamsInitBase):
    """
        Xavier 参数初始化方法
    """

    def __init__(self):
        pass

    def getParams(self, filter_size, filter_num, channels):
        weight = 1 / np.math.sqrt(filter_num * filter_size * filter_size)
        x = np.random.randn(filter_num, channels,
                            filter_size, filter_size) * weight
        b = np.zeros((filter_num))
        return x, b

    def getAffineParams(self, x_shape, output_size):
        if len(x_shape) == 4:
            shape = x_shape[1]*x_shape[2]*x_shape[3]
        else:
            shape = x_shape[1]
        weight = 1 / np.math.sqrt(shape)
        x = np.random.randn(shape, output_size) * weight
        b = np.zeros(output_size)
        return x, b
